{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Vit\n",
    "\n",
    "https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py\n",
    "\n",
    "传统卷积神经网络在提取局部特征方面表现出色，但难以捕捉图像中远距离区域之间的关系。为了解决这一问题，Dosovitskiy等人在2020年提出了Vision Transformer（ViT），首次将Transformer结构成功应用于图像分类任务。ViT的核心思想是：将图像划分为若干小块（patch），并将每个图像块视为一个“词”，送入Transformer进行全局建模。相比卷积网络，Transformer更容易捕捉全局依赖，从而提升模型对图像整体结构的理解能力。ViT在大规模数据集上取得了优异成绩，标志着视觉领域由卷积向注意力机制的转变，并深刻影响了后续视觉模型的设计。\n",
    "\n",
    "![\"tutorials\\image_src\\VIT.png\"](../../image_src/VIT.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# helpers\n",
    "def pair(t):\n",
    "    return t if isinstance(t, tuple) else (t, t)\n",
    "\n",
    "# classes\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, dim, hidden_dim, dropout = 0.):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.LayerNorm(dim),\n",
    "            nn.Linear(dim, hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(hidden_dim, dim),\n",
    "            nn.Dropout(dropout)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):\n",
    "        super().__init__()\n",
    "        inner_dim = dim_head *  heads\n",
    "        project_out = not (heads == 1 and dim_head == dim)\n",
    "\n",
    "        self.heads = heads\n",
    "        self.scale = dim_head ** -0.5\n",
    "\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "\n",
    "        self.attend = nn.Softmax(dim = -1)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)\n",
    "\n",
    "        self.to_out = nn.Sequential(\n",
    "            nn.Linear(inner_dim, dim),\n",
    "            nn.Dropout(dropout)\n",
    "        ) if project_out else nn.Identity()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.norm(x)\n",
    "\n",
    "        qkv = self.to_qkv(x).chunk(3, dim = -1)\n",
    "        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)\n",
    "\n",
    "        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale\n",
    "\n",
    "        attn = self.attend(dots)\n",
    "        attn = self.dropout(attn)\n",
    "\n",
    "        out = torch.matmul(attn, v)\n",
    "        out = rearrange(out, 'b h n d -> b n (h d)')\n",
    "        return self.to_out(out)\n",
    "\n",
    "\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):\n",
    "        super().__init__()\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "        self.layers = nn.ModuleList([])\n",
    "        for _ in range(depth):\n",
    "            self.layers.append(nn.ModuleList([\n",
    "                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),\n",
    "                FeedForward(dim, mlp_dim, dropout = dropout)\n",
    "            ]))\n",
    "\n",
    "    def forward(self, x):\n",
    "        for attn, ff in self.layers:\n",
    "            x = attn(x) + x\n",
    "            x = ff(x) + x\n",
    "\n",
    "        return self.norm(x)\n",
    "\n",
    "\n",
    "class ViT(nn.Module):\n",
    "    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):\n",
    "        super().__init__()\n",
    "        image_height, image_width = pair(image_size)\n",
    "        patch_height, patch_width = pair(patch_size)\n",
    "\n",
    "        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'\n",
    "\n",
    "        num_patches = (image_height // patch_height) * (image_width // patch_width)\n",
    "        patch_dim = channels * patch_height * patch_width\n",
    "        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'\n",
    "\n",
    "        self.to_patch_embedding = nn.Sequential(\n",
    "            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),\n",
    "            nn.LayerNorm(patch_dim),\n",
    "            nn.Linear(patch_dim, dim),\n",
    "            nn.LayerNorm(dim),\n",
    "        )\n",
    "\n",
    "        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))\n",
    "        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))\n",
    "        self.dropout = nn.Dropout(emb_dropout)\n",
    "\n",
    "        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)\n",
    "\n",
    "        self.pool = pool\n",
    "        self.to_latent = nn.Identity()\n",
    "\n",
    "        self.mlp_head = nn.Linear(dim, num_classes)\n",
    "\n",
    "    def forward(self, img):\n",
    "        x = self.to_patch_embedding(img)\n",
    "        b, n, _ = x.shape\n",
    "\n",
    "        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = b)\n",
    "        x = torch.cat((cls_tokens, x), dim=1)\n",
    "        x += self.pos_embedding[:, :(n + 1)]\n",
    "        x = self.dropout(x)\n",
    "\n",
    "        x = self.transformer(x)\n",
    "\n",
    "        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]\n",
    "\n",
    "        x = self.to_latent(x)\n",
    "        return self.mlp_head(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inps = torch.randn((2, 4, 128, 128))\n",
    "\n",
    "vit = ViT(image_size=256,\n",
    "          patch_size=2,\n",
    "          num_classes=3,\n",
    "          dim=512,\n",
    "          depth=2,\n",
    "          heads=8,\n",
    "          mlp_dim=512,\n",
    "          channels=4)\n",
    "\n",
    "out = vit(inps)\n",
    "out.size()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
